# -*- coding: utf-8 -*-
"""
Created on Tue Apr 13 08:54:49 2021
Description : CNN算法，包含训练、计算效果（pearson）、保存模型

@author: Lenovo-pc
"""

import os
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
import tensorflow as tf
import numpy as np

class cnn_1d:
    def __init__(self):
        pass
    
    def train(self,train_data,train_labels):
        assert (len(train_data)==len(train_labels)) and\
            len(train_data)>0 and len(train_labels)>0,\
        'CNN_1D ERROR:train input'

        row_size,col_size = train_data.shape
        input_shape =(row_size,int(col_size/3),3) #样本数\编码数（通道数）\特征数
        train_data = train_data.reshape(input_shape)
        
        self.model = tf.keras.models.Sequential([
            tf.keras.layers.Conv1D(filters=100, 
                                   kernel_size=(3),
                                   activation='relu',
                                    padding = 'same' ,
                                   input_shape = input_shape[1:]
                                            ),
            tf.keras.layers.MaxPooling1D(pool_size=3, padding = 'same' ),
            tf.keras.layers.Flatten(),
            tf.keras.layers.Dense(units=64,activation='relu'),
            tf.keras.layers.Dense(1)       
            ])
        
        optimizer = tf.keras.optimizers.RMSprop(0.001)
        self.model.compile(optimizer = optimizer, 
                           loss = 'mse',
                           metrics = ['mae','mse']
                           )
        self.model.fit(train_data,train_labels,epochs=5)
        
        
    
    def estimate_cor(self,test_labels,predict_test):
        cor = np.corrcoef(test_labels,predict_test)
        return cor
    
    def save_model(self,model_save_path):
        self.model.save(model_save_path)
    
    def load_model(self,model_save_path):
        self.model = tf.keras.models.load_model(model_save_path)
        
        
    def predict_data(self,input_data):
        row_size,col_size = input_data.shape
        input_shape =(row_size,int(col_size/3),3) #样本数\编码数（通道数）\特征数
        input_data = input_data.reshape(input_shape)
        pred_data = self.model.predict(input_data)
        pred_data = pred_data.reshape(len(input_data))
        return pred_data
    
if __name__ =="__main__":
    pass
    
    